# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
训练好的 Keras 模型使用、 测试结果生成

# 当前项目: Cifar10-Classification
# 创建时间: 2019/6/22 17:29
# 开发作者: vincent
# 创建平台: PyCharm Community Edition
# 版    本: V1.0
"""
import os
import cv2
import json
import time
import shutil
import numpy as np
from collections import Counter
from sklearn.metrics import classification_report, confusion_matrix
from tensorflow import keras


# 测试报告类对象
class TestReport(object):

    def __init__(self, y_ground_truth, label_mapping, digits=3, result_path=None):
        """
        :param y_ground_truth:   测试数据真实标签值
        :param label_mapping:    标签映射
        :param digits:           结果保存位数
        :param result_path:      结果保存路径
        """
        # 判断 y_ground_truth 中是否有非 label_mapping 中的类别
        assert set(y_ground_truth).issubset(set(label_mapping.keys())), 'Error: Ground_Truth Label error!!!'

        # 时区差  当前时间为北京 东八区时间  此处操作为解决Docker中以 零时区 为标准而导致的时间混乱
        time_zone = 8 - int(time.strftime('%z', time.localtime())) / 100
        # 开始时间 字符型
        self.time_str = time.strftime('%Y%m%d-%H%M%S', time.localtime(time.time() + time_zone * 60 * 60))

        # 结果保存路径，若为None 为默认为当前路径
        if result_path is None:
            self.result_path = os.path.join(os.path.dirname(__file__), '_TestResult')
        os.makedirs(self.result_path, exist_ok=True)
        
        self.label_mapping = label_mapping
        self.label_index_mapping = {_value: _key for _key, _value in label_mapping.items()}
        self.digits = digits
        self.y_ground_truth = y_ground_truth

        # 标签名称
        self.target_names = [_key for _key, _value in
                             sorted(label_mapping.items(), key=lambda x: x[1])]

        print('Label Mapping: {}'.format(sorted(self.label_mapping.items(), key=lambda x: x[1])))

    # 测试报告、混淆矩阵生成
    def __call__(self, y_predict_p):
        """
        :param y_predict_p:   模型预测结果，各个类别概率值
        :return:
        """

        # 预测标签值
        self.y_predict = [self.label_index_mapping[y_predict_i] for y_predict_i in y_predict_p.argmax(axis=1)]
        # ####################### 分类结果报告
        report = classification_report(self.y_ground_truth, self.y_predict,
                                       target_names=self.target_names,
                                       labels=self.target_names, digits=self.digits)

        report = report.split('\n')
        del report[-4]  # 删除 micro avg 行
        del report[-3]  # 删除 macro avg 行
        # 合并显示内容
        report_string = '\n'.join(report)
        print('Report: \n{}'.format(report_string))

        # ####################### 分类结果混淆矩阵
        test_confusion_matrix = confusion_matrix(self.y_ground_truth, self.y_predict,
                                                 labels=self.target_names)

        confusion_matrix_string = self._print_confusion_matrix(test_confusion_matrix)
        print('Confusion_matrix: \n{}'.format(confusion_matrix_string))

        return report_string, confusion_matrix_string

    # 混淆矩阵格式化
    def _print_confusion_matrix(self, test_confusion_matrix):
        """
        :param test_confusion_matrix:  confusion_matrix 输出混淆矩阵
        :return:
        """
        # 各个类别真实数量
        total_truth_num = test_confusion_matrix.sum(axis=1)
        # 各个类别预测数量
        total_predict_num = test_confusion_matrix.sum(axis=0)

        # 类别名称显示列格式
        target_name_len = max([len(_target) for _target in self.target_names])
        target_name_format = '{{:>{:d}s}}  '.format(target_name_len+2)
        # 各类别数量显示格式
        num_format = ' {{:>{:d}d}}'.format(len(str(len(self.y_ground_truth))))
        # 各类别索引显示格式
        label_format = ' {{:>{:d}s}}'.format(len(str(len(self.y_ground_truth))))
        # 每行汇总显示格式
        string_format = target_name_format + num_format * len(self.target_names) + '   ' + num_format

        # 显示字符内容
        print_string = list()
        # 第一行显示内容
        print_str = (target_name_format + label_format * len(self.target_names)).format(
            '', *([chr(65+i) for i in range(len(self.target_names))]))
        print_string.append(print_str)
        print_string.append('')
        # 混淆矩阵
        for i, target_i_name in enumerate(self.target_names):
            print_str = string_format.format(target_i_name, *list(test_confusion_matrix[i, :]),
                                             total_truth_num[i])
            print_string.append(print_str)
        # 最后一行显示内容
        print_str = string_format.format('', *total_predict_num, len(self.y_ground_truth))
        print_string.append('')
        print_string.append(print_str)

        # 合并显示内容
        confusion_matrix_string = '\n'.join(print_string)

        return confusion_matrix_string

    # 错误图片归类
    def error_picture_cat(self, ground_truth_path):
        """
        :param ground_truth_path:    测试图片路径列表
        :return:
        """
        # 判断测试图片路径长度是否与标签值数量一样
        assert len(ground_truth_path) == len(self.y_ground_truth), \
            'Error: ImagePath [{}] is not equal GroundTruth Label [{}]!!!'.format(
                len(ground_truth_path), len(self.y_ground_truth))

        # 判断每个测试结果，并把错误结果保存到对应文件夹
        for img_path_i, ground_truth_i, predict_i in zip(ground_truth_path, self.y_ground_truth, self.y_predict):
            if ground_truth_i != predict_i:
                save_folder_temp = os.path.join(self.result_path, self.time_str, 
                                                '{}---{}'.format(ground_truth_i, predict_i))
                os.makedirs(save_folder_temp, exist_ok=True)
                shutil.copy(img_path_i, os.path.join(save_folder_temp, os.path.basename(img_path_i)))


# 验证数据
def val_data(val_data_csv_path, input_image_shape):
    """
    :param val_data_csv_path:  验证数据 CSV 列表
    :param input_image_shape:  图片 Resize 大小
    :return:
    """
    # 验证数据列表
    val_img_path = list()
    val_img_label = list()
    for img_info in open(val_data_csv_path, 'r', encoding='utf-8'):
        img_info = img_info.strip().split(',')
        # 训练图片路径
        val_img_path.append(img_info[0])
        # 训练图片标签
        val_img_label.append(img_info[1])

    # 第一行为标签名
    val_img_path = val_img_path[1:]
    val_img_label = val_img_label[1:]

    x_val = np.zeros((len(val_img_path), input_image_shape, input_image_shape, 3), dtype='uint8')

    for i, (img_path, img_label) in enumerate(zip(val_img_path, val_img_label)):
        # 读取图片数据
        img_data = cv2.imdecode(np.fromfile(img_path, dtype='uint8'), -1)
        # 图片大小重新设定
        x_val[i] = cv2.resize(img_data, (input_image_shape, input_image_shape))

    print('ValData Label Count:{}'.format(sorted(Counter(val_img_label).items(),
                                                 key=lambda x: [x[1], x[1]])))
    return x_val, val_img_path, val_img_label


def main(test_data_path, config_path, model_path):
    """
    :param test_data_path:  测试数据 CSV 表格 路径
    :param config_path:     模型 Config 路径
    :param model_path:      模型 h5 路径
    :return:
    """

    # #################################################################
    # 读取配置信息
    config = json.load(open(config_path, 'r', encoding='utf-8'))

    # 读取测试数据
    x_val, val_img_path, y_ground_truth = val_data(test_data_path, 224)

    # ############## 导入模型
    model = keras.models.load_model(model_path)
    # ############## 模型预测
    y_predict = model.predict(x_val)

    # 打印测试报告
    report_fun = TestReport(y_ground_truth, config['label_mapping'])
    report_fun(y_predict)

    # 分类错误图片归类
    report_fun.error_picture_cat(val_img_path)

    # 写入测试结果表中
    test_result_csv = os.path.join(report_fun.result_path,
                                   '{}_{}.csv'.format(report_fun.time_str, os.path.basename(model_path)))
    with open(test_result_csv, 'w', encoding='utf-8') as f_result:
        for line_i, img_info in enumerate(open(test_data_path, 'r', encoding='utf-8')):
            img_info = img_info.strip()
            if line_i == 0:
                f_result.write('{},{}\n'.format(img_info, model_path))
            else:
                f_result.write('{},{}\n'.format(img_info, report_fun.y_predict[line_i-1]))

if __name__ == '__main__':
    test_data_path_ = r'D:\Desktop\Cifar10-Classification\Data\Image\test_data.csv'
    config_path_ = r'Z:\_temp\Cifar\__result\model.resnet_20200220-014128_ModelConfig.json'
    model_path_ = r'Z:\_temp\Cifar\__result\model.resnet_CheckPoint\model.resnet_20200220-014128_lately.h5'

    # config_path_ = r'D:\__result\model.resnet_20190623-010450_ModelConfig.json'
    # model_path_ = r'D:\__result\model.resnet_CheckPoint\model.resnet_20190623-010450_lately.h5'

    main(test_data_path_, config_path_, model_path_)
